package com.troila.rpc;

import io.grpc.Attributes;
import io.grpc.Context;
import io.grpc.Contexts;
import io.grpc.Grpc;
import io.grpc.Metadata;
import io.grpc.ServerCall;
import io.grpc.ServerCallHandler;
import io.grpc.ServerInterceptor;
import io.grpc.Status;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

@Component
@Slf4j
public class MetaDataInterceptor implements ServerInterceptor {

    private static final Metadata.Key<String> METADATA_KEY_HEADER_REAL_CLIENT =
            Metadata.Key.of("X-Forwarded-Host", Metadata.ASCII_STRING_MARSHALLER);
    private static final Metadata.Key<String> METADATA_KEY_TOKEN =
            Metadata.Key.of("token", Metadata.ASCII_STRING_MARSHALLER);

    @Override
    public <Q, P> ServerCall.Listener<Q> interceptCall(ServerCall<Q, P> call,
                                                       Metadata headers,
                                                       ServerCallHandler<Q, P> next) {
        String address = getClientAddress(call.getAttributes(), headers);
        Object token = headers.get(METADATA_KEY_TOKEN);
        log.debug("client's token is {}", token);
        if (!StringUtils.isEmpty(token) && !StringUtils.isEmpty(address)) {
            return Contexts.interceptCall(Context.current(), call, headers, next);
        }
        call.close(Status.UNAUTHENTICATED, headers);
        log.warn("error intercept unauthenticated address : {} | token : {}", address, token);
        return new ServerCall.Listener<Q>() {
        };
    }

    private String getClientAddress(Attributes attributes, Metadata metadata) {
        if (metadata.keys().contains(METADATA_KEY_HEADER_REAL_CLIENT.name())) {
            // 说明是rpc-web客户端, 经过Nginx转发
            String remoteAddr = metadata.get(METADATA_KEY_HEADER_REAL_CLIENT);
            return String.join("", "/", remoteAddr);
        }
        return attributes.get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR).toString();
    }

}
